use std::{f32::consts::PI, mem::size_of};

use bytemuck::cast_slice;
use nalgebra::{vector, Matrix4, UnitQuaternion, Vector3};
use wgpu::{include_wgsl, util::DeviceExt};

use crate::{Color, PerspectiveCamera};

mod mesh; pub use mesh::Mesh;

pub struct Renderer3D {
    pub camera: PerspectiveCamera,
    camera_proj_buffer: wgpu::Buffer,
    camera_pos_aspect_buffer: wgpu::Buffer,
    camera_bind_group: wgpu::BindGroup,
    
    simple_pipeline: wgpu::RenderPipeline,
}

impl Renderer3D {
    pub fn new(
        device: &wgpu::Device,
        format: wgpu::TextureFormat,
    ) -> Self {
        let camera = PerspectiveCamera::new(
            [-1.0, 0.0, 0.0],
            0.0, 0.0, 1.0,
            PI * 0.4, 0.01, 100.0
        );
        let camera_proj_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
            label: Some("camera_proj_buffer"),
            contents: cast_slice(&camera.calc_matrix()),
            usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
        });
        let camera_pos_aspect_buffer = device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("camera_pos_aspect_buffer"),
            size: size_of::<[f32; 4]>() as u64,
            usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });
        let camera_bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: Some("camera"),
            layout: &camera_bind_group_layout(device),
            entries: &[
                wgpu::BindGroupEntry {
                    binding: 0,
                    resource: camera_proj_buffer.as_entire_binding()
                },
                wgpu::BindGroupEntry {
                    binding: 1,
                    resource: camera_pos_aspect_buffer.as_entire_binding()
                }
            ]
        });

        let vertex_shader = include_wgsl!("../shaders/3d_vertex.wgsl");
        let vertex_module = device.create_shader_module(vertex_shader);
        let vertex_state = wgpu::VertexState {
            module: &vertex_module,
            entry_point: "vs_main",
            compilation_options: Default::default(),
            buffers: &[
                wgpu::VertexBufferLayout {
                    array_stride: size_of::<[[f32; 3]; 2]>() as u64,
                    step_mode: wgpu::VertexStepMode::Vertex,
                    attributes: &wgpu::vertex_attr_array![0 => Float32x3, 1 => Float32x3],
                }
            ]
        };

        let simple_shader = include_wgsl!("../shaders/simple3d.wgsl");
        let simple_module = device.create_shader_module(simple_shader);
        let simple_pipeline = device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {
            label: SimpleMaterial::label(),
            layout: Some(&device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
                label: SimpleMaterial::label(),
                bind_group_layouts: &[
                    &camera_bind_group_layout(device),
                    &object_bind_group_layout(device),
                    &SimpleMaterial::bind_group_layout(device)
                ],
                push_constant_ranges: &[]
            })),
            vertex: vertex_state.clone(),
            cache: Default::default(),
            primitive: Default::default(),
            depth_stencil: None,
            multisample: Default::default(),
            fragment: Some(wgpu::FragmentState {
                module: &simple_module,
                entry_point: "fs_main",
                compilation_options: Default::default(),
                targets: &[Some(wgpu::ColorTargetState {
                    format,
                    blend: Some(wgpu::BlendState::ALPHA_BLENDING),
                    write_mask: wgpu::ColorWrites::ALL,
                })]
            }),
            multiview: None,
        });

        Self {
            camera, camera_pos_aspect_buffer, camera_proj_buffer, camera_bind_group,
            simple_pipeline,
        }
    }
    pub fn update_camera(&self, queue: &wgpu::Queue) {
        let proj_matrix = self.camera.calc_matrix();
        queue.write_buffer(&self.camera_proj_buffer, 0, cast_slice(&proj_matrix));
        let pos_aspect = [self.camera.position.x, self.camera.position.y, self.camera.position.z, self.camera.aspect];
        queue.write_buffer(&self.camera_pos_aspect_buffer, 0, cast_slice(&pos_aspect));
    }
    pub fn draw<'a>(
        & self,
        transform: &Transform3D,
        mesh: &Mesh,
        pass: &mut wgpu::RenderPass,
        material: &impl Material,
    ) {
        pass.set_bind_group(0, &self.camera_bind_group, &[]);

        match material.ty() {
            MaterialType::Simple => {
                pass.set_bind_group(1, &transform.bind_group, &[]);
                material.affect_pass(pass);
                pass.set_pipeline(&self.simple_pipeline);
                pass.set_vertex_buffer(0, mesh.buffer.slice(..));
                pass.draw(0..mesh.vertices.len() as u32, 0..1);
            }
        }
    }
}

pub fn camera_bind_group_layout(device: &wgpu::Device) -> wgpu::BindGroupLayout {
    device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
        label: Some("camera"),
        entries: &[
            wgpu::BindGroupLayoutEntry {
                binding: 0,
                visibility: wgpu::ShaderStages::VERTEX,
                ty: wgpu::BindingType::Buffer {
                    ty: wgpu::BufferBindingType::Uniform,
                    has_dynamic_offset: false,
                    min_binding_size: None,
                },
                count: None,
            },
            wgpu::BindGroupLayoutEntry {
                binding: 1,
                visibility: wgpu::ShaderStages::VERTEX_FRAGMENT,
                ty: wgpu::BindingType::Buffer {
                    ty: wgpu::BufferBindingType::Uniform,
                    has_dynamic_offset: false,
                    min_binding_size: None,
                },
                count: None,
            }
        ]
    })
}

pub fn object_bind_group_layout(device: &wgpu::Device) -> wgpu::BindGroupLayout {
    device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor{
        label: Some("object"),
        entries: &[
            wgpu::BindGroupLayoutEntry {
                binding: 0,
                visibility: wgpu::ShaderStages::VERTEX,
                ty: wgpu::BindingType::Buffer {
                    ty: wgpu::BufferBindingType::Uniform,
                    has_dynamic_offset: false,
                    min_binding_size: None,
                },
                count: None,
            }
        ]
    })
}

pub struct Transform3D {
    pub position: Vector3<f32>,
    pub rotation: UnitQuaternion<f32>,
    buffer: wgpu::Buffer,
    bind_group: wgpu::BindGroup
}
impl Transform3D {
    pub fn new(device: &wgpu::Device, position: Vector3<f32>, rotation: UnitQuaternion<f32>) -> Self {
        let projection = 
            Matrix4::new_translation(&position) * Matrix4::new_rotation(rotation.scaled_axis());
        let buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
            label: Some("transform projection"),
            contents: cast_slice(projection.as_slice()),
            usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
        });
        let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: Some("transform projection"),
            layout: &object_bind_group_layout(device),
            entries: &[
                wgpu::BindGroupEntry {
                    binding: 0,
                    resource: buffer.as_entire_binding()
                }
            ]
        });
        Self {
            position, rotation, buffer, bind_group
        }
    }
    pub fn update(&self, queue: &wgpu::Queue) {
        queue.write_buffer(&self.buffer, 0, cast_slice(self.projection().as_slice()));
    }
    pub fn projection(&self) -> Matrix4<f32> {
        Matrix4::new_translation(&self.position) * Matrix4::new_rotation(self.rotation.scaled_axis())
    }
    pub fn default(device: &wgpu::Device) -> Self {
        Self::new(device, vector![0.0, 0.0, 0.0], UnitQuaternion::identity())
    }
}


pub enum MaterialType {
    Simple,
}

pub trait Material {
    fn ty(&self) -> MaterialType;
    fn affect_pass(&self, pass: &mut wgpu::RenderPass);
}


pub struct SimpleMaterial {
    pub color: Color,
    buffer: wgpu::Buffer,
    bind_group: wgpu::BindGroup,
}
impl SimpleMaterial {
    const fn label() -> Option<&'static str> {
        Some("simple material")
    }
    pub fn update(&self, queue: &wgpu::Queue) {
        queue.write_buffer(&self.buffer, 0, cast_slice(&[self.color]));
    }
    fn bind_group_layout(device: &wgpu::Device) -> wgpu::BindGroupLayout {
        device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
            label: Self::label(),
            entries: &[
                wgpu::BindGroupLayoutEntry {
                    binding: 0,
                    visibility: wgpu::ShaderStages::FRAGMENT,
                    ty: wgpu::BindingType::Buffer {
                        ty: wgpu::BufferBindingType::Uniform,
                        has_dynamic_offset: false,
                        min_binding_size: None,
                    },
                    count: None,
                }
            ]
        })
    }
    pub fn new(device: &wgpu::Device, color: Color) -> Self {
        let buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
            label: Self::label(),
            contents: &cast_slice(&[color]),
            usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::UNIFORM,
        });
        let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: Self::label(),
            layout: &Self::bind_group_layout(device),
            entries: &[
                wgpu::BindGroupEntry {
                    binding: 0,
                    resource: buffer.as_entire_binding(),
                }
            ]
        });
        Self {
            color, buffer, bind_group
        }
    }
}
impl Material for SimpleMaterial {
    fn ty(&self) -> MaterialType {
        MaterialType::Simple
    }
    fn affect_pass(&self, pass: &mut wgpu::RenderPass) {
        pass.set_bind_group(2, &self.bind_group, &[]);
    }
}